/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package bd.commons.math.distribution;

import com.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;

import commons.util.ExceptionUtil;

import java.math.BigDecimal;
import java.util.Collection;
import java.util.SortedMap;
import java.util.TreeMap;

/**
 * <a href="http://en.wikipedia.org/wiki/Sample_space">Sample space</a> with finite elements and
 * their associated probabilities. Samples can be drawn from this sample space acording to its
 * element probabilities.
 *
 * @param <E> element type
 * @author Susanta Tewari
 * @version 1.1.0
 * @history Created on 05/23/2012.
 * @since 1.1.0
 */
public class SampleSpace<E> {

    /** Field description */
    private final BiMap<Integer, E> integer_element_map;

    /** Field description */
    private final ConnectedIntegerDistribution intDist;

    /**
     * Creates a shallow copy of the passed arrays.
     *
     * @param elements all the elements in the sample space
     * @param probWeights probability weights (need not add up to {@code 1.0}) for the specified
     * elements in order;
     */
    public SampleSpace(E[] elements, BigDecimal[] probWeights) {

        // CHECK ARGS
        if (elements == null) {
            ExceptionUtil.throwIllegalArgExNull("elements");
        }

        if (probWeights == null) {
            ExceptionUtil.throwIllegalArgExNull("probs");
        }

        if (elements.length == 0) {
            throw new IllegalArgumentException("empty elements");
        }

        if (probWeights.length == 0) {
            throw new IllegalArgumentException("empty probs");
        }


        // SUM THE WEIGHTS
        BigDecimal sum = BigDecimal.ZERO;

        for (BigDecimal weight : probWeights) {

            sum = sum.add(weight);
        }


        // init map
        integer_element_map = HashBiMap.create();

        SortedMap<Integer, BigDecimal> integer_prob_weight_map = new TreeMap<>();

        for (int i = 0; i < probWeights.length; i++) {

            integer_element_map.put(i + 1, elements[i]);

            integer_prob_weight_map.put(i + 1,
                                        probWeights[i].divide(sum,
                                            Singletons.MATHCONTEXT_128_HALF_UP));

        }

        // discrete dist
        intDist = new ConnectedIntegerDistribution(integer_prob_weight_map);

    }

    /**
     * Method description
     *
     *
     * @return
     */
    public Collection<E> elements() {

        return integer_element_map.values();
    }


    /**
     * @param element an element in the sample space
     * @return probability of the specified element
     */
    public BigDecimal elementProbability(E element) {

        if (!integer_element_map.containsValue(element)) {
            throw new IllegalArgumentException("element not found");
        }

        int x = integer_element_map.inverse().get(element);

        return intDist.probabilityBD(x);
    }


    /**
     * Draws a random sample from this sample space according to its probability distribution.
     *
     * @return a random sample from this sample space
     */
    public E sample() {

        int x = intDist.sample();

        return integer_element_map.get(x);
    }
}
